from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
from PIL import Image
import torch

"""transformers 库中的 VisionEncoderDecoderModel 是一个「万能胶水」类：
把任意 Transformer 视觉编码器（ViT、Swin、BEiT…）和任意文本解码器（GPT-2、RoBERTa、MBart…）拼成一套 Encoder-Decoder 架构，
专门做「看图说话」式的生成任务——图像描述、OCR、文档解析等都能用"""
model = VisionEncoderDecoderModel.from_pretrained("D:/ideaSpace/MyPython/models/vit-gpt2-image-captioning")
processor = ViTImageProcessor.from_pretrained("D:/ideaSpace/MyPython/models/vit-gpt2-image-captioning")
tokenizer = AutoTokenizer.from_pretrained("D:/ideaSpace/MyPython/models/vit-gpt2-image-captioning")

# 处理png格式图片会报错
# image = Image.open(r"D:\ideaSpace\rag-in-action-master\wukong.jpg")
image = Image.open(r"D:\ideaSpace\ai_home_work\kecheng\jiagou.jpg")
pixel_values = processor(images=image, return_tensors="pt").pixel_values

with torch.no_grad():
    # generated = model.generate(pixel_values, max_length=30, num_beams=4)
    generated = model.generate(
        pixel_values,
        max_length=30,
        do_sample=True,      # 或 False 用贪心
        top_p=0.9,
        temperature=0.7
    )
caption = tokenizer.batch_decode(generated, skip_special_tokens=True)[0]
print("输出结果：", caption)